#include <iostream>
#include <torch/torch.h>
#include "inculde/dataSet.h"
#include "inculde/model.h"
using namespace std;

int main(){
    // 真实方程 y = 3 * x + 1;

    Model model = Model();

    int length = dataSet.sizes()[0];
    int flag = 0;
    int k = 0;
    while(true){
        for (int i = 0; i < length; i++){
            k++;
            // 1. 将数据逐个放入
            model.forward(dataSet[i][0]);
            // 2. 计算损失值
            // model.loss_function(dataSet[i][1]);
            // 3. 反向传播
            flag = model.backward(dataSet[i][1]);
            if (flag == 1)
                break;
        }
        if (flag == 1)
            break;
    }
    std::cout << "***********************\n";
    std::cout << "***********************\n";
    cout << "stop training\n";
    cout << "final loss: " << model.loss << endl;
    cout << "final weight: " << model.weight << endl;
    cout << "final bias: " << model.bias << endl;
    cout << "training times: " << k << endl;
    std::cout << "***********************\n";
    std::cout << "***********************\n";

    auto test = model.forward(torch::tensor({20,}));
    cout << "test : " << test.data() << endl;  // expect 61

    return 0;
}

/*
 * 必须动态的修改学习速率，否则会让参数错过最低点
 * */
